import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.dates import DateFormatter
import yfinance as yf

# 获取股票信息
data = yf.download("BTC-USD", period="3mo")
ohlc_data = data[["Open", "High", "Low", "Close"]]


# MACD参数设置
def calculate_macd(data, fast=12, slow=26, signal=9):
    data["EMA12"] = data["Close"].ewm(span=fast).mean()
    data["EMA26"] = data["Close"].ewm(span=slow).mean()
    data["MACD"] = data["EMA12"] - data["EMA26"]
    data["Signal"] = data["MACD"].ewm(span=signal).mean()
    return data

# KDJ计算
def calculate_kdj(data, n=9):
    min_low = data["Low"].rolling(n).min()
    max_high = data["High"].rolling(n).max()
    data["RSV"] = (data["Close"] - min_low) / (max_high - min_low) * 100
    data["K"] = data["RSV"].ewm(com=2).mean()
    data["D"] = data["K"].ewm(com=2).mean()
    data["J"] = 3 * data["K"] - 2 * data["D"]
    return data

# 交易信号生成
def generate_signals(data):
    # MACD金叉/死叉
    data["MACD_Cross"] = np.where(data["MACD"] > data["Signal"], 1, -1)

    # KDJ超买超卖
    data["KDJ_Buy"] = (data["J"] < 20) & (data["J"].shift(1) >= 20)
    data["KDJ_Sell"] = (data["J"] > 80) & (data["J"].shift(1) <= 80)

    # 综合信号
    data["Signal"] = np.select(
        [
            (data["MACD_Cross"] == 1) & data["KDJ_Buy"],
            (data["MACD_Cross"] == -1) & data["KDJ_Sell"],
        ],
        [1, -1],
        default=0,
    )
    return data

# 可视化视图
def plot_backtest(data):
    plt.figure(figsize=(16, 12))
    ax1 = plt.subplot(3, 1, 1)
    ax2 = plt.subplot(3, 1, 2, sharex=ax1)
    ax3 = plt.subplot(3, 1, 3, sharex=ax1)

    # 价格与交易信号
    ax1.plot(data["Close"], label="Price", color="#1f77b4")
    buy_signals = data[data["Signal"] == 1]
    sell_signals = data[data["Signal"] == -1]
    ax1.scatter(
        buy_signals.index,
        buy_signals["Close"],
        marker="^",
        color="g",
        s=100,
        label="Buy",
    )
    ax1.scatter(
        sell_signals.index,
        sell_signals["Close"],
        marker="v",
        color="r",
        s=100,
        label="Sell",
    )

    # MACD可视化
    ax2.plot(data["MACD"], label="MACD", color="#ff7f0e")
    ax2.plot(data["Signal"], label="Signal", color="#2ca02c")
    ax2.bar(
        data.index,
        data["MACD"] - data["Signal"],
        color=np.where(data["MACD"] > data["Signal"], "lime", "red"),
    )

    # KDJ可视化
    ax3.plot(data["K"], label="K", color="#d62728")
    ax3.plot(data["D"], label="D", color="#9467bd")
    ax3.plot(data["J"], label="J", color="#8c564b", linestyle="--")
    ax3.axhline(80, color="gray", linestyle=":")
    ax3.axhline(20, color="gray", linestyle=":")

    # 美化设置
    ax1.set_title("MACD+KDJ Trading Strategy Backtest", fontsize=14)
    ax2.set_title("MACD Indicator", fontsize=12)
    ax3.set_title("KDJ Indicator", fontsize=12)
    plt.xticks(rotation=45)
    ax1.legend()
    plt.tight_layout()
    plt.show()


# 回测性能分析
def performance_analysis(data):
    data["Return"] = data["Close"].pct_change()
    data["Strategy_Return"] = data["Signal"].shift(1) * data["Return"]
    data["Cumulative_Return"] = (1 + data["Strategy_Return"]).cumprod()

    print(f"累计收益率: {data['Cumulative_Return'][-1]:.2%}")
    print(
        f"最大回撤: {(1 - data['Cumulative_Return'] / data['Cumulative_Return'].cummax()).max():.2%}"
    )

if __name__ == "__main__":
    ohlc_data = calculate_macd(ohlc_data)
    ohlc_data = calculate_kdj(ohlc_data)
    ohlc_data = generate_signals(ohlc_data)
    plot_backtest(ohlc_data)
    performance_analysis(ohlc_data)
